# coding:utf-8
import pymongo
import redis
import hashlib
import six

from dotenv import load_dotenv
from pathlib import Path

env_path = Path('') / '.env'
load_dotenv(dotenv_path=env_path)


def init_db():
    my_client = pymongo.MongoClient('localhost',27017)  # 连接地址
    my_db = my_client['spider']  # 数据库名
    return my_db


class Redis(object):
    def __init__(self):
        self.redis_host = '127.0.0.1'
        self.redis_port = 6379
        self.redis_db = 0
        self.storage = self._get_storage()
        self.hash_func = getattr(hashlib, 'md5')

    def _get_hash_value(self, data):
        """
        根据给定的数据，返回的对应信息摘要hash值
        :param data: 给定的原始数据（二进制类型的字符串数据）
        :return: hash值
        """
        hash_obj = self.hash_func()
        hash_obj.update(self._safe_data(data))  # python3   bytes    python2 str
        hash_value = hash_obj.hexdigest()
        return hash_value

    @staticmethod
    def _safe_data(data):
        """
        python2   str  === python3   bytes
        python2   uniocde === python3  str
        :param data: 给定的原始数据
        :return: 二进制类型的字符串数据
        """
        if six.PY3:
            if isinstance(data, bytes):
                return data
            elif isinstance(data, str):
                return data.encode()
            else:
                raise Exception("请提供一个字符串")  # 建议使用英文来描述
        else:
            if isinstance(data, str):
                return data
            elif isinstance(data, unicode):
                return data.encode()
            else:
                raise Exception("请提供一个字符串")  # 建议使用英文来描述

    def _get_storage(self):
        """返回一个redis连接对象"""
        pool = redis.ConnectionPool(host=self.redis_host, port=self.redis_port, db=self.redis_db)
        client = redis.StrictRedis(connection_pool=pool)
        return client

    def save(self, redis_key, data):
        """
        根据data计算出对应的指纹进行存储
        :param redis_key:
        :param data: 给定的原始数据
        :return: 存储的结果
        """
        hash_value = self._get_hash_value(data)
        return self.storage.sadd(redis_key, hash_value)

    def is_exists(self, redis_key, data):
        """
        判断给定的数据对应的指纹是否存在
        :param redis_key:
        :param data: 给定的原始数据
        :return: True or False
        """
        hash_value = self._get_hash_value(data)
        return self.storage.sismember(redis_key, hash_value)

